package com.ting.test.algorithms.堆.加强堆;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;

public class Heap<T> {
    ArrayList<T> heap = new ArrayList();
    int heapSize = 0;
    Map<T, Integer> indexMap = new HashMap<T, Integer>();
    Comparator<? super T> comparator;

    public Heap(Comparator<? super T> comparator) {
        this.comparator = comparator;
    }

    /**
     * 弹出指定元素
     *
     * @param t
     * @return
     */
    public int heapPop(T t) {
        if (!indexMap.containsKey(t)) {
            return -1;
        }
        Integer index = indexMap.get(t);
        bigPop(index);
        return index;
    }

    private void bigPop(Integer index) {
        swap(index, --heapSize);
        indexMap.remove(heap.get(heapSize));
        heap.remove(index);
        int left = 2 * index + 1;
        int largest;
        while (left < heapSize) {
            //如果没有右子节点
            if (left + 1 >= heapSize) {
                largest = left;
            } else {//如果有右子节点
                largest = comparator.compare(heap.get(left), heap.get(left + 1)) > 0 ? left : left + 1;
            }


            if (comparator.compare(heap.get(index), heap.get(largest)) > 0) {
                break;
            }
            swap(index, largest);
            index = largest;
            left = 2 * index + 1;
        }
    }

    /**
     * 返回下标
     *
     * @param t
     * @return
     */
    public int heapInsert(T t) {
        //先判断元素是否存在
        //存在则直接覆盖
        if (indexMap.containsKey(t)) {
            return indexMap.get(t);
        }

        heap.add(t);
        int index = bigInsert(heapSize);
        indexMap.put(t, index);
        heapSize++;
        return index;

    }

    private int bigInsert(int index) {
        while (comparator.compare(heap.get(index), heap.get((index - 1) / 2)) > 0) {
            swap(index, (index - 1) / 2);
            index = (index - 1) / 2;
        }
        return index;
    }

    private void swap(int index, int index2) {
        T temp = heap.get(index);
        heap.set(index, heap.get(index2));
        heap.set(index2, temp);
        indexMap.put(heap.get(index), index);
        indexMap.put(heap.get(index2), index2);
    }
}
